import torch
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import ImageFolder
from model import GoogLeNet, Inception
from torch import nn
import pandas as pd
import time
from PIL import Image

if (__name__ == "__main__"):
    model = GoogLeNet(Inception)
    model.load_state_dict(torch.load("./model/best_model.pth"))
    image = Image.open("d1.jpg")

    normalize = transforms.Normalize([0.162, 0.151, 0.138], [0.058, 0.052, 0.048])
    # 定义数据集处理方法变量
    test_transforms = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor(), normalize])
    image = test_transforms(image)
    # 添加批次维度
    image = image.unsqueeze(0)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    classes_names = ['猫', '狗']
    with torch.no_grad():
        model.eval()
        image = image.to(device)
        output = model(image)
        pre_lab = torch.argmax(output, dim=1)
        result = pre_lab.item()
    print(pre_lab)
    print(classes_names[result])
